# main.py
import os
import shutil
import time

from flask import Flask, Response, render_template, url_for, request, jsonify
from flask_cors import CORS
import threading
import cv2
import numpy as np
from werkzeug.utils import secure_filename, redirect
from threading import Thread

from detect import VideoTracker
from shared import shared_data, shared_data_lock, SharedData  # 👈 确保这样导入
from utils.parser import get_config
from yolo import YOLO

app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})



shared_data = SharedData()

UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'mp4', 'avi', 'mov'}

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
uploaded_video_path = None
video_processor = None
current_video_source = None
processing_lock = threading.Lock()

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@app.route('/upload_video', methods=['POST'])
def upload_video():
    global uploaded_video_path

    if 'video' not in request.files:
        return jsonify({"status": "error", "message": "No video part"}), 400

    file = request.files['video']

    if file.filename == '':
        return jsonify({"status": "error", "message": "Empty filename"}), 400

    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)

        os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
        file.save(filepath)

        uploaded_video_path = filepath
        print(f"[DEBUG] Uploaded video path set to: {filepath}")

        try:
            global video_processor
            video_processor = VideoProcessor(config, filepath)
            print("[INFO] VideoProcessor initialized successfully.")
        except Exception as e:
            print(f"[ERROR] Failed to initialize VideoProcessor: {e}")
            return Response(f"Error initializing video processor: {e}", status=500)

        return jsonify({"status": "success"})
    else:
        return jsonify({"status": "error", "message": "Invalid file type"}), 400


class Config:
    def __init__(self):
        self.YOLOWEIGHT = "./weights/yolov5s.pt"
        self.CARSPEEDLIMIT = 25
        self.TRAFFICJAMLIMIT = 16
        self.SAVEVIDEO = True
        self.CARCLASSIFY = True
        self.forwardline = []
        self.leftline = []
        self.rightline = []
        self.trafficline1 = []
        self.trafficline2 = []
        self.laneline = []
        self.leftlight = []
        self.forwardlight = []
        self.rightlight = []
        self.preflag = [0, 0, 0, 0, 0, 0]


config = Config()


class VideoProcessor:
    def __init__(self, config, video_source=0):
        self.config = config
        self.cfg = get_config()
        self.cfg.merge_from_file("./configs/deep_sort.yaml")
        self.video_tracker = VideoTracker(self.cfg)
        self.video_tracker.detector = YOLO(config.YOLOWEIGHT)

        self.video_source = video_source
        self.cap = None
        self.running = True

        self.output_dir = os.path.join(os.getcwd(), "output", "camera_0")
        if os.path.exists(self.output_dir):
            for filename in os.listdir(self.output_dir):
                file_path = os.path.join(self.output_dir, filename)
                try:
                    if os.path.isfile(file_path) or os.path.islink(file_path):
                        os.unlink(file_path)
                    elif os.path.isdir(file_path):
                        if filename == "illegal":
                            continue
                        shutil.rmtree(file_path)
                except Exception as e:
                    print(f'Failed to delete {file_path}. Reason: {e}')
        os.makedirs(self.output_dir, exist_ok=True)

        self.video_tracker.outputDir = self.output_dir

    def stop(self):
        self.running = False
        if self.cap and self.cap.isOpened():  # 先检查 cap 是否已初始化且打开
            self.cap.release()

    def get_frame(self):
        if not self.running:
            return None

        ret, frame = self.cap.read()
        if not ret:
            print("⚠️ 读取视频帧失败，可能是文件结束或损坏。")
            return None
        return frame if ret else None

    def process_frame(self):
        if self.cap is None:
            self.cap = cv2.VideoCapture(self.video_source)
            if not self.cap.isOpened():
                raise ValueError(f"❌ 无法打开视频源: {self.video_source}")
            print(f"✅ 成功打开视频源: {self.video_source}")
            self.video_tracker._init_for_stream(self.config)  # 👈 确保初始化一次

        ret, frame = self.cap.read()
        if not ret:
            print("⚠️ 读取视频帧失败，可能是文件结束或损坏。")
            return None

        processed_frame = self.video_tracker.process_frame(frame, self.config)
        if processed_frame is None:
            print("⚠️ processed_frame 为 None")
            return None

        with shared_data_lock:
            shared_data.frame = processed_frame.copy()

        return processed_frame


# 在 gen 函数中加入帧率控制
def gen(video_processor):
    print("[DEBUG] Starting video generation loop...")
    try:
        while video_processor.running:
            frame = video_processor.process_frame()
            if frame is not None:
                print("[DEBUG] Got a valid frame")
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                ret, jpeg = cv2.imencode('.jpg', frame)
                if ret:
                    yield (
                        b'--frame\r\n'
                        b'Content-Type: image/jpeg\r\n\r\n' + jpeg.tobytes() + b'\r\n'
                    )
            else:
                print("[ERROR] Frame is None, waiting...")
            time.sleep(0.03)
    except Exception as e:
        print(f"[CRITICAL] Error in gen(): {e}")








@app.route('/')
def index():
    return render_template('index.html')


# 在 app.py 中顶部定义全局变量
video_processor = None
current_video_source = None
processing_lock = threading.Lock()


# 修改 gen 函数
def gen(video_processor):
    try:
        while video_processor.running:
            frame = video_processor.process_frame()
            if frame is not None:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                ret, jpeg = cv2.imencode('.jpg', frame)
                if ret:
                    yield (
                        b'--frame\r\n'
                        b'Content-Type: image/jpeg\r\n\r\n' + jpeg.tobytes() + b'\r\n'
                    )
                else:
                    print("[ERROR] Failed to encode frame")
            else:
                print("[DEBUG] Frame is None, waiting...")
            time.sleep(0.03)  # 避免 CPU 占满
    except Exception as e:
        print(f"[CRITICAL] Error in gen(): {e}")
        import traceback
        traceback.print_exc()
    finally:
        print("[INFO] Gen function exited")


# 修改视频流路由
@app.route('/camera_feed')
def camera_feed():
    global video_processor, current_video_source

    # 如果当前没有视频处理器或当前不是摄像头源
    if video_processor is None or current_video_source != 'camera':
        # 创建新的 VideoProcessor 实例
        with processing_lock:
            if video_processor is not None:
                video_processor.stop()

            video_processor = VideoProcessor(config, 0)  # 0 是默认摄像头
            current_video_source = 'camera'

    return Response(gen(video_processor),
                    mimetype='multipart/x-mixed-replace; boundary=frame')


# 添加一个新的视频文件流路由
@app.route('/video_feed')
def video_feed():
    global video_processor, current_video_source, uploaded_video_path

    print("[DEBUG] 当前路径:", os.getcwd())
    print("[DEBUG] 视频路径:", uploaded_video_path)

    if uploaded_video_path is None:
        return Response("No video uploaded", status=404)

    if video_processor is None or current_video_source != 'video':
        with processing_lock:
            if video_processor is not None:
                video_processor.stop()
            print("[DEBUG] 创建新的 VideoProcessor 实例")
            video_processor = VideoProcessor(config, uploaded_video_path)
            current_video_source = 'video'

    return Response(gen(video_processor),
                    mimetype='multipart/x-mixed-replace; boundary=frame')



if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, threaded=True, debug=False)
